// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT

#include "gemm_quantization_common.hpp"
#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_wmma_cshuffle_v3.hpp"

namespace ck {
namespace tensor_operation {
namespace device {
namespace instance {

template <typename OutElementOp,
          BlockGemmPipelineScheduler GemmPipelineScheduler,
          BlockGemmPipelineVersion GemmPipeline>
using device_gemm_quantization_wmma_c_shuffle_i8_i8_i8_km_kn_mn_instances = std::tuple<
    // clang-format off
        //################################| ALayout| BLayout|   DsLayout| ELayout|  AData|  BData|      DsData|    EData| AccData| CShuffle|           A|           B|           CDE|   GemmSpec| Block|  MPer|  NPer|  KPer| AK1| BK1| MPer| NPer| MRepeat| NRepeat|    ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|    BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds|   CShuffle|   CShuffle| CShuffleBlockTransfer| CDEShuffleBlockTransfer|               BlkGemm|     BlkGemm| ComputeTypeA| ComputeTypeB|
        //################################|        |        |           |        |   Type|   Type|        Type|     Type|    Type| DataType| Elementwise| Elementwise|   Elementwise|           |  Size| Block| Block| Block|    |    | Wmma| Wmma|        |        |     ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar|    ExtraM|     ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar|    ExtraN|    MRepeat|    NRepeat|        ClusterLengths|        ScalarPerVectors|             PipeSched| PipelineVer|             |             |
        //################################|        |        |           |        |       |       |            |         |        |         |   Operation|   Operation|     Operation|           |      |      |      |      |    |    |     |     |        |        | Lengths_AK0_M_AK1|   ArrangeOrder|               |               |      PerVector|  PerVector_AK1|          | Lengths_BK0_N_BK1|   ArrangeOrder|               |               |      PerVector|  PerVector_BK1|          | PerShuffle| PerShuffle|     _MBlock_MPerBlock|                        |                      |            |             |             |
        //################################|        |        |           |        |       |       |            |         |        |         |            |            |              |           |      |      |      |      |    |    |     |     |        |        |                  |               |               |               |               |               |          |                  |               |               |               |               |               |          |           |           |     _NBlock_NPerBlock|                        |                      |            |             |             |
        DeviceGemmMultipleD_Wmma_CShuffleV3<    Col,    Row, Empty_Tuple,     Row, int8_t, int8_t, Empty_Tuple,   int8_t, int32_t,  int32_t, PassThrough, PassThrough,  OutElementOp, MNKPadding,   256,   128,   128,    64,   8,   8,   16,   16,       4,       2,       S<4, 64, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,      true,       S<4, 64, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,      true,          1,          1,       S<1, 32, 1, 8>,                     S<1>, GemmPipelineScheduler, GemmPipeline,      int8_t,      int8_t>,
        DeviceGemmMultipleD_Wmma_CShuffleV3<    Col,    Row, Empty_Tuple,     Row, int8_t, int8_t, Empty_Tuple,   int8_t, int32_t,  int32_t, PassThrough, PassThrough,  OutElementOp, MNKPadding,   256,   128,   128,    64,   8,   8,   16,   16,       4,       2,       S<4, 32, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,      true,       S<4, 32, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,      true,          1,          1,       S<1, 32, 1, 8>,                     S<1>, GemmPipelineScheduler, GemmPipeline,      int8_t,      int8_t>,
        DeviceGemmMultipleD_Wmma_CShuffleV3<    Col,    Row, Empty_Tuple,     Row, int8_t, int8_t, Empty_Tuple,   int8_t, int32_t,  int32_t, PassThrough, PassThrough,  OutElementOp, MNKPadding,   128,   128,   128,    32,   8,   8,   16,   16,       4,       4,       S<4, 32, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,      true,       S<4, 32, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,      true,          1,          1,       S<1, 32, 1, 4>,                     S<1>, GemmPipelineScheduler, GemmPipeline,      int8_t,      int8_t>,
        DeviceGemmMultipleD_Wmma_CShuffleV3<    Col,    Row, Empty_Tuple,     Row, int8_t, int8_t, Empty_Tuple,   int8_t, int32_t,  int32_t, PassThrough, PassThrough,  OutElementOp, MNKPadding,    64,    32,    64,    64,   8,   8,   16,   16,       2,       2,       S<4, 16, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,     false,       S<4, 16, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,     false,          1,          1,       S<1, 16, 1, 4>,                     S<1>, GemmPipelineScheduler, GemmPipeline,      int8_t,      int8_t>
    // clang-format on
    >;

template <typename OutElementOp,
          BlockGemmPipelineScheduler GemmPipelineScheduler,
          BlockGemmPipelineVersion GemmPipeline>
using device_gemm_quantization_wmma_c_shuffle_i8_i8_i8_km_nk_mn_instances = std::tuple<
    // clang-format off
        //################################| ALayout| BLayout|   DsLayout| ELayout| AData|  BData|      DsData|    EData| AccData| CShuffle|           A|           B|           CDE|   GemmSpec| Block|  MPer|  NPer|  KPer| AK1| BK1| MPer| NPer| MRepeat| NRepeat|    ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|    BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds|   CShuffle|   CShuffle| CShuffleBlockTransfer| CDEShuffleBlockTransfer|               BlkGemm|     BlkGemm| ComputeTypeA| ComputeTypeB|
        //################################|        |        |           |        |  Type|   Type|        Type|     Type|    Type| DataType| Elementwise| Elementwise|   Elementwise|           |  Size| Block| Block| Block|    |    | Wmma| Wmma|        |        |     ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar|    ExtraM|     ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar|    ExtraN|    MRepeat|    NRepeat|        ClusterLengths|        ScalarPerVectors|             PipeSched| PipelineVer|             |             |
        //################################|        |        |           |        |      |       |            |         |        |         |   Operation|   Operation|     Operation|           |      |      |      |      |    |    |     |     |        |        | Lengths_AK0_M_AK1|   ArrangeOrder|               |               |      PerVector|  PerVector_AK1|          | Lengths_BK0_N_BK1|   ArrangeOrder|               |               |      PerVector|  PerVector_BK1|          | PerShuffle| PerShuffle|     _MBlock_MPerBlock|                        |                      |            |             |             |
        //################################|        |        |           |        |      |       |            |         |        |         |            |            |              |           |      |      |      |      |    |    |     |     |        |        |                  |               |               |               |               |               |          |                  |               |               |               |               |               |          |           |           |     _NBlock_NPerBlock|                        |                      |            |             |             |
        DeviceGemmMultipleD_Wmma_CShuffleV3<    Col,    Col, Empty_Tuple,    Row, int8_t, int8_t, Empty_Tuple,   int8_t, int32_t,  int32_t, PassThrough, PassThrough,  OutElementOp, MNKPadding,   256,   128,   128,    32,   8,   8,   16,   16,       4,       2,       S<4, 64, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,      true,       S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,          1,          1,       S<1, 32, 1, 8>,                     S<1>, GemmPipelineScheduler, GemmPipeline,      int8_t,      int8_t>,
        DeviceGemmMultipleD_Wmma_CShuffleV3<    Col,    Col, Empty_Tuple,    Row, int8_t, int8_t, Empty_Tuple,   int8_t, int32_t,  int32_t, PassThrough, PassThrough,  OutElementOp, MNKPadding,   256,   128,   256,    64,   8,   8,   16,   16,       4,       4,       S<4, 64, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,      true,       S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,          1,          1,       S<1, 32, 1, 8>,                     S<1>, GemmPipelineScheduler, GemmPipeline,      int8_t,      int8_t>,
        DeviceGemmMultipleD_Wmma_CShuffleV3<    Col,    Col, Empty_Tuple,    Row, int8_t, int8_t, Empty_Tuple,   int8_t, int32_t,  int32_t, PassThrough, PassThrough,  OutElementOp, MNKPadding,   256,   128,   160,    64,   8,   8,   16,   16,       2,       5,       S<4, 64, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,      true,       S<8, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,          1,          1,       S<1, 64, 1, 4>,                     S<1>, GemmPipelineScheduler, GemmPipeline,      int8_t,      int8_t>,
        DeviceGemmMultipleD_Wmma_CShuffleV3<    Col,    Col, Empty_Tuple,    Row, int8_t, int8_t, Empty_Tuple,   int8_t, int32_t,  int32_t, PassThrough, PassThrough,  OutElementOp, MNKPadding,   128,    64,    80,    64,   8,   8,   16,   16,       1,       5,       S<4, 32, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,     false,       S<8, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,     false,          1,          1,       S<1, 64, 1, 2>,                     S<1>, GemmPipelineScheduler, GemmPipeline,      int8_t,      int8_t>
    // clang-format on
    >;

template <typename OutElementOp,
          BlockGemmPipelineScheduler GemmPipelineScheduler,
          BlockGemmPipelineVersion GemmPipeline>
using device_gemm_quantization_wmma_c_shuffle_i8_i8_i8_mk_kn_mn_instances = std::tuple<
    // clang-format off
        //################################| ALayout| BLayout|   DsLayout| ELayout|  AData|  BData|      DsData|    EData| AccData| CShuffle|           A|           B|           CDE|   GemmSpec| Block|  MPer|  NPer|  KPer| AK1| BK1| MPer| NPer| MRepeat| NRepeat|    ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|    BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds|   CShuffle|   CShuffle| CShuffleBlockTransfer| CDEShuffleBlockTransfer|               BlkGemm|     BlkGemm| ComputeTypeA| ComputeTypeB|
        //################################|        |        |           |        |   Type|   Type|        Type|     Type|    Type| DataType| Elementwise| Elementwise|   Elementwise|           |  Size| Block| Block| Block|    |    | Wmma| Wmma|        |        |     ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar|    ExtraM|     ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar|    ExtraN|    MRepeat|    NRepeat|        ClusterLengths|        ScalarPerVectors|             PipeSched| PipelineVer|             |             |
        //################################|        |        |           |        |       |       |            |         |        |         |   Operation|   Operation|     Operation|           |      |      |      |      |    |    |     |     |        |        | Lengths_AK0_M_AK1|   ArrangeOrder|               |               |      PerVector|  PerVector_AK1|          | Lengths_BK0_N_BK1|   ArrangeOrder|               |               |      PerVector|  PerVector_BK1|          | PerShuffle| PerShuffle|     _MBlock_MPerBlock|                        |                      |            |             |             |
        //################################|        |        |           |        |       |       |            |         |        |         |            |            |              |           |      |      |      |      |    |    |     |     |        |        |                  |               |               |               |               |               |          |                  |               |               |               |               |               |          |           |           |     _NBlock_NPerBlock|                        |                      |            |             |             |
        DeviceGemmMultipleD_Wmma_CShuffleV3<    Row,    Row, Empty_Tuple,     Row, int8_t, int8_t, Empty_Tuple,   int8_t, int32_t,  int32_t, PassThrough, PassThrough,  OutElementOp, MNKPadding,   256,   128,   128,    32,   8,   8,   16,   16,       4,       2,       S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,     false,       S<4, 64, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,     false,          1,          1,       S<1, 32, 1, 8>,                     S<1>, GemmPipelineScheduler, GemmPipeline,      int8_t,      int8_t>,
        DeviceGemmMultipleD_Wmma_CShuffleV3<    Row,    Row, Empty_Tuple,     Row, int8_t, int8_t, Empty_Tuple,   int8_t, int32_t,  int32_t, PassThrough, PassThrough,  OutElementOp, MNKPadding,   128,   128,    64,    64,   8,   8,   16,   16,       4,       2,       S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,     false,       S<4, 32, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,     false,          1,          1,       S<1, 32, 1, 4>,                     S<1>, GemmPipelineScheduler, GemmPipeline,      int8_t,      int8_t>,
        DeviceGemmMultipleD_Wmma_CShuffleV3<    Row,    Row, Empty_Tuple,     Row, int8_t, int8_t, Empty_Tuple,   int8_t, int32_t,  int32_t, PassThrough, PassThrough,  OutElementOp, MNKPadding,   256,   128,   256,    64,   8,   8,   16,   16,       4,       4,       S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,       S<4, 64, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,      true,          1,          1,       S<1, 32, 1, 8>,                     S<1>, GemmPipelineScheduler, GemmPipeline,      int8_t,      int8_t>,
        DeviceGemmMultipleD_Wmma_CShuffleV3<    Row,    Row, Empty_Tuple,     Row, int8_t, int8_t, Empty_Tuple,   int8_t, int32_t,  int32_t, PassThrough, PassThrough,  OutElementOp, MNKPadding,   256,   128,   128,    64,   8,   8,   16,   16,       4,       2,       S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,       S<4, 64, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,      true,          1,          1,       S<1, 32, 1, 8>,                     S<1>, GemmPipelineScheduler, GemmPipeline,      int8_t,      int8_t>
    // clang-format on
    >;

template <typename OutElementOp,
          BlockGemmPipelineScheduler GemmPipelineScheduler,
          BlockGemmPipelineVersion GemmPipeline>
using device_gemm_quantization_wmma_c_shuffle_i8_i8_i8_mk_nk_mn_instances = std::tuple<
    // clang-format off
        //################################| ALayout| BLayout|   DsLayout| ELayout|  AData|  BData|      DsData|    EData| AccData| CShuffle|           A|           B|           CDE|   GemmSpec| Block|  MPer|  NPer|  KPer| AK1| BK1| MPer| NPer| MRepeat| NRepeat|    ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|    BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds|   CShuffle|   CShuffle| CShuffleBlockTransfer| CDEShuffleBlockTransfer|               BlkGemm|     BlkGemm| ComputeTypeA| ComputeTypeB|
        //################################|        |        |           |        |   Type|   Type|        Type|     Type|    Type| DataType| Elementwise| Elementwise|   Elementwise|           |  Size| Block| Block| Block|    |    | Wmma| Wmma|        |        |     ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar|    ExtraM|     ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar|    ExtraN|    MRepeat|    NRepeat|        ClusterLengths|        ScalarPerVectors|             PipeSched| PipelineVer|             |             |
        //################################|        |        |           |        |       |       |            |         |        |         |   Operation|   Operation|     Operation|           |      |      |      |      |    |    |     |     |        |        | Lengths_AK0_M_AK1|   ArrangeOrder|               |               |      PerVector|  PerVector_AK1|          | Lengths_BK0_N_BK1|   ArrangeOrder|               |               |      PerVector|  PerVector_BK1|          | PerShuffle| PerShuffle|     _MBlock_MPerBlock|                        |                      |            |             |             |
        //################################|        |        |           |        |       |       |            |         |        |         |            |            |              |           |      |      |      |      |    |    |     |     |        |        |                  |               |               |               |               |               |          |                  |               |               |               |               |               |          |           |           |     _NBlock_NPerBlock|                        |                      |            |             |             |
        DeviceGemmMultipleD_Wmma_CShuffleV3<    Row,    Col, Empty_Tuple,     Row, int8_t, int8_t, Empty_Tuple,   int8_t, int32_t,  int32_t, PassThrough, PassThrough,  OutElementOp, MNKPadding,   256,   128,   128,    32,   8,   8,   16,   16,       4,       2,       S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,       S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,          1,          1,        S<1, 32, 1, 8>,                    S<1>, GemmPipelineScheduler, GemmPipeline,      int8_t,      int8_t>,
        DeviceGemmMultipleD_Wmma_CShuffleV3<    Row,    Col, Empty_Tuple,     Row, int8_t, int8_t, Empty_Tuple,   int8_t, int32_t,  int32_t, PassThrough, PassThrough,  OutElementOp, MNKPadding,   128,   128,    64,    64,   8,   8,   16,   16,       4,       2,       S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,     false,       S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,     false,          1,          1,        S<1, 32, 1, 4>,                    S<1>, GemmPipelineScheduler, GemmPipeline,      int8_t,      int8_t>,
        DeviceGemmMultipleD_Wmma_CShuffleV3<    Row,    Col, Empty_Tuple,     Row, int8_t, int8_t, Empty_Tuple,   int8_t, int32_t,  int32_t, PassThrough, PassThrough,  OutElementOp, MNKPadding,   256,   128,   256,    64,   8,   8,   16,   16,       4,       4,       S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,       S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,          1,          1,        S<1, 32, 1, 8>,                    S<1>, GemmPipelineScheduler, GemmPipeline,      int8_t,      int8_t>,
        DeviceGemmMultipleD_Wmma_CShuffleV3<    Row,    Col, Empty_Tuple,     Row, int8_t, int8_t, Empty_Tuple,   int8_t, int32_t,  int32_t, PassThrough, PassThrough,  OutElementOp, MNKPadding,    64,    32,    64,    64,   8,   8,   16,   16,       2,       2,       S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,     false,       S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,     false,          1,          1,        S<1, 16, 1, 4>,                    S<1>, GemmPipelineScheduler, GemmPipeline,      int8_t,      int8_t>
    // clang-format on
    >;

} // namespace instance
} // namespace device
} // namespace tensor_operation
} // namespace ck
